
import re

def diagnose_reward_fn(data_source, solution_str, ground_truth, extra_info=None):
    """
    医学多模态诊断奖励函数
    ======================
    综合判定模型输出格式和诊断内容。
    
    评分规则：
    1. 格式分（最高0.2）：
        - <thinking>、<answer>标签完整闭合
        - <region>与<description>成对出现，region坐标合法
        - 至少有一组region/description
        - answer标签内有内容
    2. 内容分（最高0.8）：
        - ground_truth为标准答案（多个疾病用逗号分隔）
        - answer内容（逗号分割strip后）完全覆盖所有标准答案：0.8
        - 只覆盖部分：按命中比例给分（如1/2疾病：0.4）
        - answer中如包含误诊（非标准答案内容），每个误诊扣0.1（下限0.1）
    3. 总分：格式分+内容分。格式不合格时内容分最多给0.1（鼓励输出），格式合格内容全对才满分。
    """

    # -------- 1. 检查格式 --------
    fmt_score = 0
    content_score = 0
    min_score = 0.1  # 最低分

    # 检查思维链标签
    has_thinking = re.search(r'<thinking>.*</thinking>', solution_str, re.DOTALL)
    has_answer = re.search(r'<answer>.*</answer>', solution_str, re.DOTALL)

    # region/description成对、坐标检查
    region_desc_pairs = re.findall(
        r'<region>\((\d+),\s*(\d+),\s*(\d+),\s*(\d+)\)</region>\s*<description>(.*?)</description>',
        solution_str, re.DOTALL)

    valid_coords = True
    for x1, y1, x2, y2, _ in region_desc_pairs:
        coords = [int(x1), int(y1), int(x2), int(y2)]
        if not all(0 <= v <= 224 for v in coords):
            valid_coords = False
            break

    # 至少一组region/description、标签成对、坐标合法
    if has_thinking and has_answer and region_desc_pairs and valid_coords:
        fmt_score = 0.2
    else:
        fmt_score = 0.0

    # -------- 2. 检查诊断内容 --------
    # 只从answer标签内提取内容
    answer_match = re.search(r'<answer>(.*?)</answer>', solution_str, re.DOTALL)
    if answer_match:
        answer_text = answer_match.group(1)
        # split多个疾病，strip空格
        pred_labels = [x.strip() for x in re.split(r'[，,]\s*', answer_text) if x.strip()]
    else:
        pred_labels = []

    # 处理标准答案
    gold_labels = [x.strip() for x in ground_truth.split(", ") if x.strip()]

    # 统计命中
    hit = sum(1 for l in gold_labels if any(l in pred for pred in pred_labels))
    miss = len(gold_labels) - hit

    # 统计误诊
    wrong = [pred for pred in pred_labels if not any(g in pred for g in gold_labels)]
    wrong_cnt = len(wrong)

    # 内容给分
    if gold_labels:
        if hit == 0:
            content_score = 0.0
        else:
            content_score = 0.8 * (hit / len(gold_labels))
            # 误诊扣分，每个误诊扣0.5，最低扣到0.1,因为到这里的判别说明至少有一个输出是正确的，为了鼓励输出
            content_score = max(content_score - 0.5 * wrong_cnt, min_score)
    else:
        # 没有标准答案（极少见），不给分
        content_score = 0.0

    # -------- 3. 合并得分 --------
    if fmt_score == 0.2:
        total_score = fmt_score + content_score  # 满分1.0
        total_score = min(total_score, 1.0)
    else:
        # 格式不对，内容最多0.1
        total_score = min(fmt_score + min(content_score, 0.1), 0.2)

    # 强制分数范围
    total_score = max(min_score, min(total_score, 1.0))
    return total_score
